# -*- coding: utf-8 -*-
"""
Created on Sat Feb  1 16:13:12 2020

@author: lenovo03
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import talib

#本地读取数据
def ReadStockData():
    df = pd.read_csv('603688-2019.csv')
    df = df[['trade_date','close']]
    return df
df = ReadStockData()

price = df['close']
date = df['trade_date']

#CUSUM计算逻辑
def detect_via_cusum_lg(ts, istart,threshold_times):
    """
    detect a time series using  cusum algorithm
    :param ts: the time series to be detected
    :param istart: the data from index 0 to index istart will be used as cold startup data to train 启动数据
    :param threshold_times: the times for setting threshold  设置阈值的时间
    :return:
    """
    
    S_h = 0  #high
    S_l = 0  #low
    S_list = np.zeros(istart)  #设置前istart个值为0[0., 0., 0.]

    meanArray = talib.SMA(ts,timeperiod = istart)  #244,[nan,3.6755,3.6625,...]
    '''
    简介：简单移动平均线SMA,talib输入数据必须是ndarray类型
    用法：talib.SMA(close, timeperiod=30)
    返回值：一位数组（numpy.ndarray）
    ''' 
    #np.log()自然对数
    stdArray = talib.STDDEV(np.log(ts/meanArray),timeperiod = istart)  #标准差
    #print(stdArray)
    '''
    简介：标准差
    用法：talib.STDDEV(close.timeperiod=5,nbdev=1)
    '''
    for i in range(istart, len(ts)):     #2----243
        tslog = np.log(ts[i] / meanArray[i - 1])
        print("{}tslog值:{}".format(date[i],tslog))
        
        S_h_ = max(0, S_h + tslog - stdArray[i-1])
        S_l_ = min(0, S_l + tslog + stdArray[i-1])

        if S_h_> threshold_times*stdArray[i-1]:
            S_list = np.append(S_list,1)
            #print("{}:上涨信号".format(date_list[i]))
            S_h_ = 0
        elif abs(S_l_)> threshold_times*stdArray[i-1]:
            S_list = np.append(S_list, -1)
            #print("{}:下跌信号".format(date_list[i]))
            S_l_ = 0
        else:
            S_list = np.append(S_list, 0)
            #print("{}:无信号".format(date_list[i]))
        S_h = S_h_
        S_l = S_l_
    return S_list

up_signal_list,down_signal_list = [],[]#分别装上涨、下跌信号
signal_list = detect_via_cusum_lg(price,istart=2,threshold_times=9)  #244=242+2
#如[0,0,1,-1,...,0]，装每个交易日的信号，无信号0，上涨1，下跌-1
print(len(signal_list))

total_up,total_down = 0,0

for i in range(0,len(signal_list)):
    if signal_list[i] == 1:
        #up_signal_list标明第i个交易日有上涨信号，用于画图
        up_signal_list.append(i)
    elif signal_list[i] == -1 :
        #down_signal_list标明第i个交易日有下跌信号，用于画图
        down_signal_list.append(i)

print("上涨信号总计{}个".format(len(up_signal_list)))
print("下跌信号总计{}个".format(len(down_signal_list)))

plt.figure(figsize=(10,5))
plt.plot(price, color='y', lw=2.)
#plt.title(r'CUSUM',fontproperties='SimHei',fontsize=20)
plt.xlabel('日期',fontproperties='SimHei',fontsize=15)
plt.ylabel('价格',fontproperties='SimHei',fontsize=15)
plt.plot(price, '^', markersize=2, color='r', label='UP signal', markevery=up_signal_list)
plt.plot(price, 'v', markersize=2, color='g', label='DOWN signal', markevery=down_signal_list)
plt.legend()
plt.grid(True)
plt.show()